{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TextCNN\n",
    "\n",
    "TextCNN利用CNN（卷积神经网络）进行文本特征抽取，不同大小的卷积核分别抽取n-gram特征，卷积计算出的特征图经过MaxPooling保留最大的特征值，然后将拼接成一个向量作为文本的表示。\n",
    "\n",
    "这里我们基于TextCNN原始论文的设定，分别采用了100个大小为2,3,4的卷积核，最后得到的文本向量大小为100*3=300维。\n",
    "\n",
    "![TextCNN](img/cnn.jpeg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-07-17 11:37:20,835 INFO: Use cuda: True, gpu id: 0.\n"
     ]
    }
   ],
   "source": [
    "import logging\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format='%(asctime)-15s %(levelname)s: %(message)s')\n",
    "\n",
    "# set seed\n",
    "seed = 666\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "# set cuda\n",
    "gpu = 0\n",
    "use_cuda = gpu >= 0 and torch.cuda.is_available()\n",
    "if use_cuda:\n",
    "    torch.cuda.set_device(gpu)\n",
    "    device = torch.device(\"cuda\", gpu)\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "logging.info(\"Use cuda: %s, gpu id: %d.\", use_cuda, gpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-07-17 11:37:25,526 INFO: Fold lens [1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000]\n"
     ]
    }
   ],
   "source": [
    "# split data to 10 fold\n",
    "fold_num = 10\n",
    "data_file = '../data/train_set.csv'\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "def all_data2fold(fold_num, num=10000):\n",
    "    fold_data = []\n",
    "    f = pd.read_csv(data_file, sep='\\t', encoding='UTF-8')\n",
    "    texts = f['text'].tolist()[:num]\n",
    "    labels = f['label'].tolist()[:num]\n",
    "\n",
    "    total = len(labels)\n",
    "\n",
    "    index = list(range(total))\n",
    "    np.random.shuffle(index)\n",
    "\n",
    "    all_texts = []\n",
    "    all_labels = []\n",
    "    for i in index:\n",
    "        all_texts.append(texts[i])\n",
    "        all_labels.append(labels[i])\n",
    "\n",
    "    label2id = {}\n",
    "    for i in range(total):\n",
    "        label = str(all_labels[i])\n",
    "        if label not in label2id:\n",
    "            label2id[label] = [i]\n",
    "        else:\n",
    "            label2id[label].append(i)\n",
    "\n",
    "    all_index = [[] for _ in range(fold_num)]\n",
    "    for label, data in label2id.items():\n",
    "        # print(label, len(data))\n",
    "        batch_size = int(len(data) / fold_num)\n",
    "        other = len(data) - batch_size * fold_num\n",
    "        for i in range(fold_num):\n",
    "            cur_batch_size = batch_size + 1 if i < other else batch_size\n",
    "            # print(cur_batch_size)\n",
    "            batch_data = [data[i * batch_size + b] for b in range(cur_batch_size)]\n",
    "            all_index[i].extend(batch_data)\n",
    "\n",
    "    batch_size = int(total / fold_num)\n",
    "    other_texts = []\n",
    "    other_labels = []\n",
    "    other_num = 0\n",
    "    start = 0\n",
    "    for fold in range(fold_num):\n",
    "        num = len(all_index[fold])\n",
    "        texts = [all_texts[i] for i in all_index[fold]]\n",
    "        labels = [all_labels[i] for i in all_index[fold]]\n",
    "\n",
    "        if num > batch_size:\n",
    "            fold_texts = texts[:batch_size]\n",
    "            other_texts.extend(texts[batch_size:])\n",
    "            fold_labels = labels[:batch_size]\n",
    "            other_labels.extend(labels[batch_size:])\n",
    "            other_num += num - batch_size\n",
    "        elif num < batch_size:\n",
    "            end = start + batch_size - num\n",
    "            fold_texts = texts + other_texts[start: end]\n",
    "            fold_labels = labels + other_labels[start: end]\n",
    "            start = end\n",
    "        else:\n",
    "            fold_texts = texts\n",
    "            fold_labels = labels\n",
    "\n",
    "        assert batch_size == len(fold_labels)\n",
    "\n",
    "        # shuffle\n",
    "        index = list(range(batch_size))\n",
    "        np.random.shuffle(index)\n",
    "\n",
    "        shuffle_fold_texts = []\n",
    "        shuffle_fold_labels = []\n",
    "        for i in index:\n",
    "            shuffle_fold_texts.append(fold_texts[i])\n",
    "            shuffle_fold_labels.append(fold_labels[i])\n",
    "\n",
    "        data = {'label': shuffle_fold_labels, 'text': shuffle_fold_texts}\n",
    "        fold_data.append(data)\n",
    "\n",
    "    logging.info(\"Fold lens %s\", str([len(data['label']) for data in fold_data]))\n",
    "\n",
    "    return fold_data\n",
    "\n",
    "\n",
    "fold_data = all_data2fold(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build train, dev, test data\n",
    "fold_id = 9\n",
    "\n",
    "# dev\n",
    "dev_data = fold_data[fold_id]\n",
    "\n",
    "# train\n",
    "train_texts = []\n",
    "train_labels = []\n",
    "for i in range(0, fold_id):\n",
    "    data = fold_data[i]\n",
    "    train_texts.extend(data['text'])\n",
    "    train_labels.extend(data['label'])\n",
    "\n",
    "train_data = {'label': train_labels, 'text': train_texts}\n",
    "\n",
    "# test\n",
    "test_data_file = '../data/test_a.csv'\n",
    "f = pd.read_csv(test_data_file, sep='\\t', encoding='UTF-8')\n",
    "texts = f['text'].tolist()\n",
    "test_data = {'label': [0] * len(texts), 'text': texts}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-07-17 11:37:26,603 INFO: PyTorch version 1.2.0 available.\n",
      "2020-07-17 11:37:29,673 INFO: Build vocab: words 4337, labels 14.\n"
     ]
    }
   ],
   "source": [
    "# build vocab\n",
    "from collections import Counter\n",
    "from transformers import BasicTokenizer\n",
    "\n",
    "basic_tokenizer = BasicTokenizer()\n",
    "\n",
    "\n",
    "class Vocab():\n",
    "    def __init__(self, train_data):\n",
    "        self.min_count = 5\n",
    "        self.pad = 0\n",
    "        self.unk = 1\n",
    "        self._id2word = ['[PAD]', '[UNK]']\n",
    "        self._id2extword = ['[PAD]', '[UNK]']\n",
    "\n",
    "        self._id2label = []\n",
    "        self.target_names = []\n",
    "\n",
    "        self.build_vocab(train_data)\n",
    "\n",
    "        reverse = lambda x: dict(zip(x, range(len(x))))\n",
    "        self._word2id = reverse(self._id2word)\n",
    "        self._label2id = reverse(self._id2label)\n",
    "\n",
    "        logging.info(\"Build vocab: words %d, labels %d.\" % (self.word_size, self.label_size))\n",
    "\n",
    "    def build_vocab(self, data):\n",
    "        self.word_counter = Counter()\n",
    "\n",
    "        for text in data['text']:\n",
    "            words = text.split()\n",
    "            for word in words:\n",
    "                self.word_counter[word] += 1\n",
    "\n",
    "        for word, count in self.word_counter.most_common():\n",
    "            if count >= self.min_count:\n",
    "                self._id2word.append(word)\n",
    "\n",
    "        label2name = {0: '科技', 1: '股票', 2: '体育', 3: '娱乐', 4: '时政', 5: '社会', 6: '教育', 7: '财经',\n",
    "                      8: '家居', 9: '游戏', 10: '房产', 11: '时尚', 12: '彩票', 13: '星座'}\n",
    "\n",
    "        self.label_counter = Counter(data['label'])\n",
    "\n",
    "        for label in range(len(self.label_counter)):\n",
    "            count = self.label_counter[label]\n",
    "            self._id2label.append(label)\n",
    "            self.target_names.append(label2name[label])\n",
    "\n",
    "    def load_pretrained_embs(self, embfile):\n",
    "        with open(embfile, encoding='utf-8') as f:\n",
    "            lines = f.readlines()\n",
    "            items = lines[0].split()\n",
    "            word_count, embedding_dim = int(items[0]), int(items[1])\n",
    "\n",
    "        index = len(self._id2extword)\n",
    "        embeddings = np.zeros((word_count + index, embedding_dim))\n",
    "        for line in lines[1:]:\n",
    "            values = line.split()\n",
    "            self._id2extword.append(values[0])\n",
    "            vector = np.array(values[1:], dtype='float64')\n",
    "            embeddings[self.unk] += vector\n",
    "            embeddings[index] = vector\n",
    "            index += 1\n",
    "\n",
    "        embeddings[self.unk] = embeddings[self.unk] / word_count\n",
    "        embeddings = embeddings / np.std(embeddings)\n",
    "\n",
    "        reverse = lambda x: dict(zip(x, range(len(x))))\n",
    "        self._extword2id = reverse(self._id2extword)\n",
    "\n",
    "        assert len(set(self._id2extword)) == len(self._id2extword)\n",
    "\n",
    "        return embeddings\n",
    "\n",
    "    def word2id(self, xs):\n",
    "        if isinstance(xs, list):\n",
    "            return [self._word2id.get(x, self.unk) for x in xs]\n",
    "        return self._word2id.get(xs, self.unk)\n",
    "\n",
    "    def extword2id(self, xs):\n",
    "        if isinstance(xs, list):\n",
    "            return [self._extword2id.get(x, self.unk) for x in xs]\n",
    "        return self._extword2id.get(xs, self.unk)\n",
    "\n",
    "    def label2id(self, xs):\n",
    "        if isinstance(xs, list):\n",
    "            return [self._label2id.get(x, self.unk) for x in xs]\n",
    "        return self._label2id.get(xs, self.unk)\n",
    "\n",
    "    @property\n",
    "    def word_size(self):\n",
    "        return len(self._id2word)\n",
    "\n",
    "    @property\n",
    "    def extword_size(self):\n",
    "        return len(self._id2extword)\n",
    "\n",
    "    @property\n",
    "    def label_size(self):\n",
    "        return len(self._id2label)\n",
    "\n",
    "\n",
    "vocab = Vocab(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build module\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, hidden_size):\n",
    "        super(Attention, self).__init__()\n",
    "        self.weight = nn.Parameter(torch.Tensor(hidden_size, hidden_size))\n",
    "        self.weight.data.normal_(mean=0.0, std=0.05)\n",
    "\n",
    "        self.bias = nn.Parameter(torch.Tensor(hidden_size))\n",
    "        b = np.zeros(hidden_size, dtype=np.float32)\n",
    "        self.bias.data.copy_(torch.from_numpy(b))\n",
    "\n",
    "        self.query = nn.Parameter(torch.Tensor(hidden_size))\n",
    "        self.query.data.normal_(mean=0.0, std=0.05)\n",
    "\n",
    "    def forward(self, batch_hidden, batch_masks):\n",
    "        # batch_hidden: b x len x hidden_size (2 * hidden_size of lstm)\n",
    "        # batch_masks:  b x len\n",
    "\n",
    "        # linear\n",
    "        key = torch.matmul(batch_hidden, self.weight) + self.bias  # b x len x hidden\n",
    "\n",
    "        # compute attention\n",
    "        outputs = torch.matmul(key, self.query)  # b x len\n",
    "\n",
    "        masked_outputs = outputs.masked_fill((1 - batch_masks).bool(), float(-1e32))\n",
    "\n",
    "        attn_scores = F.softmax(masked_outputs, dim=1)  # b x len\n",
    "\n",
    "        # 对于全零向量，-1e32的结果为 1/len, -inf为nan, 额外补0\n",
    "        masked_attn_scores = attn_scores.masked_fill((1 - batch_masks).bool(), 0.0)\n",
    "\n",
    "        # sum weighted sources\n",
    "        batch_outputs = torch.bmm(masked_attn_scores.unsqueeze(1), key).squeeze(1)  # b x hidden\n",
    "\n",
    "        return batch_outputs, attn_scores\n",
    "\n",
    "\n",
    "# build word encoder\n",
    "word2vec_path = '../emb/word2vec.txt'\n",
    "dropout = 0.15\n",
    "\n",
    "\n",
    "class WordCNNEncoder(nn.Module):\n",
    "    def __init__(self, vocab):\n",
    "        super(WordCNNEncoder, self).__init__()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.word_dims = 100\n",
    "\n",
    "        self.word_embed = nn.Embedding(vocab.word_size, self.word_dims, padding_idx=0)\n",
    "\n",
    "        extword_embed = vocab.load_pretrained_embs(word2vec_path)\n",
    "        extword_size, word_dims = extword_embed.shape\n",
    "        logging.info(\"Load extword embed: words %d, dims %d.\" % (extword_size, word_dims))\n",
    "\n",
    "        self.extword_embed = nn.Embedding(extword_size, word_dims, padding_idx=0)\n",
    "        self.extword_embed.weight.data.copy_(torch.from_numpy(extword_embed))\n",
    "        self.extword_embed.weight.requires_grad = False\n",
    "\n",
    "        input_size = self.word_dims\n",
    "\n",
    "        self.filter_sizes = [2, 3, 4]  # n-gram window\n",
    "        self.out_channel = 100\n",
    "        self.convs = nn.ModuleList([nn.Conv2d(1, self.out_channel, (filter_size, input_size), bias=True)\n",
    "                                    for filter_size in self.filter_sizes])\n",
    "\n",
    "    def forward(self, word_ids, extword_ids):\n",
    "        # word_ids: sen_num x sent_len\n",
    "        # extword_ids: sen_num x sent_len\n",
    "        # batch_masks: sen_num x sent_len\n",
    "        sen_num, sent_len = word_ids.shape\n",
    "\n",
    "        word_embed = self.word_embed(word_ids)  # sen_num x sent_len x 100\n",
    "        extword_embed = self.extword_embed(extword_ids)\n",
    "        batch_embed = word_embed + extword_embed\n",
    "\n",
    "        if self.training:\n",
    "            batch_embed = self.dropout(batch_embed)\n",
    "\n",
    "        batch_embed.unsqueeze_(1)  # sen_num x 1 x sent_len x 100\n",
    "\n",
    "        pooled_outputs = []\n",
    "        for i in range(len(self.filter_sizes)):\n",
    "            filter_height = sent_len - self.filter_sizes[i] + 1\n",
    "            conv = self.convs[i](batch_embed)\n",
    "            hidden = F.relu(conv)  # sen_num x out_channel x filter_height x 1\n",
    "\n",
    "            mp = nn.MaxPool2d((filter_height, 1))  # (filter_height, filter_width)\n",
    "            pooled = mp(hidden).reshape(sen_num,\n",
    "                                        self.out_channel)  # sen_num x out_channel x 1 x 1 -> sen_num x out_channel\n",
    "\n",
    "            pooled_outputs.append(pooled)\n",
    "\n",
    "        reps = torch.cat(pooled_outputs, dim=1)  # sen_num x total_out_channel\n",
    "\n",
    "        if self.training:\n",
    "            reps = self.dropout(reps)\n",
    "\n",
    "        return reps\n",
    "\n",
    "\n",
    "# build sent encoder\n",
    "sent_hidden_size = 256\n",
    "sent_num_layers = 2\n",
    "\n",
    "\n",
    "class SentEncoder(nn.Module):\n",
    "    def __init__(self, sent_rep_size):\n",
    "        super(SentEncoder, self).__init__()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        self.sent_lstm = nn.LSTM(\n",
    "            input_size=sent_rep_size,\n",
    "            hidden_size=sent_hidden_size,\n",
    "            num_layers=sent_num_layers,\n",
    "            batch_first=True,\n",
    "            bidirectional=True\n",
    "        )\n",
    "\n",
    "    def forward(self, sent_reps, sent_masks):\n",
    "        # sent_reps:  b x doc_len x sent_rep_size\n",
    "        # sent_masks: b x doc_len\n",
    "\n",
    "        sent_hiddens, _ = self.sent_lstm(sent_reps)  # b x doc_len x hidden*2\n",
    "        sent_hiddens = sent_hiddens * sent_masks.unsqueeze(2)\n",
    "\n",
    "        if self.training:\n",
    "            sent_hiddens = self.dropout(sent_hiddens)\n",
    "\n",
    "        return sent_hiddens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-07-17 11:37:29,942 INFO: Load extword embed: words 5978, dims 100.\n",
      "2020-07-17 11:37:32,640 INFO: Build model with cnn word encoder, lstm sent encoder.\n",
      "2020-07-17 11:37:32,641 INFO: Model param num: 4.11 M.\n"
     ]
    }
   ],
   "source": [
    "# build model\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, vocab):\n",
    "        super(Model, self).__init__()\n",
    "        self.sent_rep_size = 300\n",
    "        self.doc_rep_size = sent_hidden_size * 2\n",
    "        self.all_parameters = {}\n",
    "        parameters = []\n",
    "        self.word_encoder = WordCNNEncoder(vocab)\n",
    "        parameters.extend(list(filter(lambda p: p.requires_grad, self.word_encoder.parameters())))\n",
    "\n",
    "        self.sent_encoder = SentEncoder(self.sent_rep_size)\n",
    "        self.sent_attention = Attention(self.doc_rep_size)\n",
    "        parameters.extend(list(filter(lambda p: p.requires_grad, self.sent_encoder.parameters())))\n",
    "        parameters.extend(list(filter(lambda p: p.requires_grad, self.sent_attention.parameters())))\n",
    "\n",
    "        self.out = nn.Linear(self.doc_rep_size, vocab.label_size, bias=True)\n",
    "        parameters.extend(list(filter(lambda p: p.requires_grad, self.out.parameters())))\n",
    "\n",
    "        if use_cuda:\n",
    "            self.to(device)\n",
    "\n",
    "        if len(parameters) > 0:\n",
    "            self.all_parameters[\"basic_parameters\"] = parameters\n",
    "\n",
    "        logging.info('Build model with cnn word encoder, lstm sent encoder.')\n",
    "\n",
    "        para_num = sum([np.prod(list(p.size())) for p in self.parameters()])\n",
    "        logging.info('Model param num: %.2f M.' % (para_num / 1e6))\n",
    "\n",
    "    def forward(self, batch_inputs):\n",
    "        # batch_inputs(batch_inputs1, batch_inputs2): b x doc_len x sent_len\n",
    "        # batch_masks : b x doc_len x sent_len\n",
    "        batch_inputs1, batch_inputs2, batch_masks = batch_inputs\n",
    "        batch_size, max_doc_len, max_sent_len = batch_inputs1.shape[0], batch_inputs1.shape[1], batch_inputs1.shape[2]\n",
    "        batch_inputs1 = batch_inputs1.view(batch_size * max_doc_len, max_sent_len)  # sen_num x sent_len\n",
    "        batch_inputs2 = batch_inputs2.view(batch_size * max_doc_len, max_sent_len)  # sen_num x sent_len\n",
    "        batch_masks = batch_masks.view(batch_size * max_doc_len, max_sent_len)  # sen_num x sent_len\n",
    "\n",
    "        sent_reps = self.word_encoder(batch_inputs1, batch_inputs2)  # sen_num x sent_rep_size\n",
    "\n",
    "        sent_reps = sent_reps.view(batch_size, max_doc_len, self.sent_rep_size)  # b x doc_len x sent_rep_size\n",
    "        batch_masks = batch_masks.view(batch_size, max_doc_len, max_sent_len)  # b x doc_len x max_sent_len\n",
    "        sent_masks = batch_masks.bool().any(2).float()  # b x doc_len\n",
    "\n",
    "        sent_hiddens = self.sent_encoder(sent_reps, sent_masks)  # b x doc_len x doc_rep_size\n",
    "        doc_reps, atten_scores = self.sent_attention(sent_hiddens, sent_masks)  # b x doc_rep_size\n",
    "\n",
    "        batch_outputs = self.out(doc_reps)  # b x num_labels\n",
    "\n",
    "        return batch_outputs\n",
    "\n",
    "\n",
    "model = Model(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build optimizer\n",
    "learning_rate = 2e-4\n",
    "decay = .75\n",
    "decay_step = 1000\n",
    "\n",
    "\n",
    "class Optimizer:\n",
    "    def __init__(self, model_parameters):\n",
    "        self.all_params = []\n",
    "        self.optims = []\n",
    "        self.schedulers = []\n",
    "\n",
    "        for name, parameters in model_parameters.items():\n",
    "            if name.startswith(\"basic\"):\n",
    "                optim = torch.optim.Adam(parameters, lr=learning_rate)\n",
    "                self.optims.append(optim)\n",
    "\n",
    "                l = lambda step: decay ** (step // decay_step)\n",
    "                scheduler = torch.optim.lr_scheduler.LambdaLR(optim, lr_lambda=l)\n",
    "                self.schedulers.append(scheduler)\n",
    "                self.all_params.extend(parameters)\n",
    "\n",
    "            else:\n",
    "                Exception(\"no nameed parameters.\")\n",
    "\n",
    "        self.num = len(self.optims)\n",
    "\n",
    "    def step(self):\n",
    "        for optim, scheduler in zip(self.optims, self.schedulers):\n",
    "            optim.step()\n",
    "            scheduler.step()\n",
    "            optim.zero_grad()\n",
    "\n",
    "    def zero_grad(self):\n",
    "        for optim in self.optims:\n",
    "            optim.zero_grad()\n",
    "\n",
    "    def get_lr(self):\n",
    "        lrs = tuple(map(lambda x: x.get_lr()[-1], self.schedulers))\n",
    "        lr = ' %.5f' * self.num\n",
    "        res = lr % lrs\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build dataset\n",
    "def sentence_split(text, vocab, max_sent_len=256, max_segment=16):\n",
    "    words = text.strip().split()\n",
    "    document_len = len(words)\n",
    "\n",
    "    index = list(range(0, document_len, max_sent_len))\n",
    "    index.append(document_len)\n",
    "\n",
    "    segments = []\n",
    "    for i in range(len(index) - 1):\n",
    "        segment = words[index[i]: index[i + 1]]\n",
    "        assert len(segment) > 0\n",
    "        segment = [word if word in vocab._id2word else '<UNK>' for word in segment]\n",
    "        segments.append([len(segment), segment])\n",
    "\n",
    "    assert len(segments) > 0\n",
    "    if len(segments) > max_segment:\n",
    "        segment_ = int(max_segment / 2)\n",
    "        return segments[:segment_] + segments[-segment_:]\n",
    "    else:\n",
    "        return segments\n",
    "\n",
    "\n",
    "def get_examples(data, vocab, max_sent_len=256, max_segment=8):\n",
    "    label2id = vocab.label2id\n",
    "    examples = []\n",
    "\n",
    "    for text, label in zip(data['text'], data['label']):\n",
    "        # label\n",
    "        id = label2id(label)\n",
    "\n",
    "        # words\n",
    "        sents_words = sentence_split(text, vocab, max_sent_len, max_segment)\n",
    "        doc = []\n",
    "        for sent_len, sent_words in sents_words:\n",
    "            word_ids = vocab.word2id(sent_words)\n",
    "            extword_ids = vocab.extword2id(sent_words)\n",
    "            doc.append([sent_len, word_ids, extword_ids])\n",
    "        examples.append([id, len(doc), doc])\n",
    "\n",
    "    logging.info('Total %d docs.' % len(examples))\n",
    "    return examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build loader\n",
    "\n",
    "def batch_slice(data, batch_size):\n",
    "    batch_num = int(np.ceil(len(data) / float(batch_size)))\n",
    "    for i in range(batch_num):\n",
    "        cur_batch_size = batch_size if i < batch_num - 1 else len(data) - batch_size * i\n",
    "        docs = [data[i * batch_size + b] for b in range(cur_batch_size)]\n",
    "\n",
    "        yield docs\n",
    "\n",
    "\n",
    "def data_iter(data, batch_size, shuffle=True, noise=1.0):\n",
    "    \"\"\"\n",
    "    randomly permute data, then sort by source length, and partition into batches\n",
    "    ensure that the length of  sentences in each batch\n",
    "    \"\"\"\n",
    "\n",
    "    batched_data = []\n",
    "    if shuffle:\n",
    "        np.random.shuffle(data)\n",
    "\n",
    "    lengths = [example[1] for example in data]\n",
    "    noisy_lengths = [- (l + np.random.uniform(- noise, noise)) for l in lengths]\n",
    "    sorted_indices = np.argsort(noisy_lengths).tolist()\n",
    "    sorted_data = [data[i] for i in sorted_indices]\n",
    "\n",
    "    batched_data.extend(list(batch_slice(sorted_data, batch_size)))\n",
    "\n",
    "    if shuffle:\n",
    "        np.random.shuffle(batched_data)\n",
    "\n",
    "    for batch in batched_data:\n",
    "        yield batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some function\n",
    "from sklearn.metrics import f1_score, precision_score, recall_score\n",
    "\n",
    "\n",
    "def get_score(y_ture, y_pred):\n",
    "    y_ture = np.array(y_ture)\n",
    "    y_pred = np.array(y_pred)\n",
    "    f1 = f1_score(y_ture, y_pred, average='macro') * 100\n",
    "    p = precision_score(y_ture, y_pred, average='macro') * 100\n",
    "    r = recall_score(y_ture, y_pred, average='macro') * 100\n",
    "\n",
    "    return str((reformat(p, 2), reformat(r, 2), reformat(f1, 2))), reformat(f1, 2)\n",
    "\n",
    "\n",
    "def reformat(num, n):\n",
    "    return float(format(num, '0.' + str(n) + 'f'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build trainer\n",
    "\n",
    "import time\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "clip = 5.0\n",
    "epochs = 1\n",
    "early_stops = 3\n",
    "log_interval = 50\n",
    "\n",
    "test_batch_size = 128\n",
    "train_batch_size = 128\n",
    "\n",
    "save_model = './cnn.bin'\n",
    "save_test = './cnn.csv'\n",
    "\n",
    "class Trainer():\n",
    "    def __init__(self, model, vocab):\n",
    "        self.model = model\n",
    "        self.report = True\n",
    "\n",
    "        self.train_data = get_examples(train_data, vocab)\n",
    "        self.batch_num = int(np.ceil(len(self.train_data) / float(train_batch_size)))\n",
    "        self.dev_data = get_examples(dev_data, vocab)\n",
    "\n",
    "        # criterion\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "        # label name\n",
    "        self.target_names = vocab.target_names\n",
    "\n",
    "        # optimizer\n",
    "        self.optimizer = Optimizer(model.all_parameters)\n",
    "\n",
    "        # count\n",
    "        self.step = 0\n",
    "        self.early_stop = -1\n",
    "        self.best_train_f1, self.best_dev_f1 = 0, 0\n",
    "        self.last_epoch = epochs\n",
    "\n",
    "    def train(self):\n",
    "        logging.info('Start training...')\n",
    "        for epoch in range(1, epochs + 1):\n",
    "            train_f1 = self._train(epoch)\n",
    "\n",
    "            dev_f1 = self._eval(epoch)\n",
    "\n",
    "            if self.best_dev_f1 <= dev_f1:\n",
    "                logging.info(\n",
    "                    \"Exceed history dev = %.2f, current dev = %.2f\" % (self.best_dev_f1, dev_f1))\n",
    "                torch.save(self.model.state_dict(), save_model)\n",
    "\n",
    "                self.best_train_f1 = train_f1\n",
    "                self.best_dev_f1 = dev_f1\n",
    "                self.early_stop = 0\n",
    "            else:\n",
    "                self.early_stop += 1\n",
    "                if self.early_stop == early_stops:\n",
    "                    logging.info(\n",
    "                        \"Eearly stop in epoch %d, best train: %.2f, dev: %.2f\" % (\n",
    "                            epoch - early_stops, self.best_train_f1, self.best_dev_f1))\n",
    "                    self.last_epoch = epoch\n",
    "                    break\n",
    "\n",
    "    def test(self):\n",
    "        self.model.load_state_dict(torch.load(save_model))\n",
    "        self._eval(self.last_epoch + 1, test=True)\n",
    "    \n",
    "    def _train(self, epoch):\n",
    "        self.optimizer.zero_grad()\n",
    "        self.model.train()\n",
    "\n",
    "        start_time = time.time()\n",
    "        epoch_start_time = time.time()\n",
    "        overall_losses = 0\n",
    "        losses = 0\n",
    "        batch_idx = 1\n",
    "        y_pred = []\n",
    "        y_true = []\n",
    "        for batch_data in data_iter(self.train_data, train_batch_size, shuffle=True):\n",
    "            torch.cuda.empty_cache()\n",
    "            batch_inputs, batch_labels = self.batch2tensor(batch_data)\n",
    "            batch_outputs = self.model(batch_inputs)\n",
    "            loss = self.criterion(batch_outputs, batch_labels)\n",
    "            loss.backward()\n",
    "\n",
    "            loss_value = loss.detach().cpu().item()\n",
    "            losses += loss_value\n",
    "            overall_losses += loss_value\n",
    "\n",
    "            y_pred.extend(torch.max(batch_outputs, dim=1)[1].cpu().numpy().tolist())\n",
    "            y_true.extend(batch_labels.cpu().numpy().tolist())\n",
    "\n",
    "            nn.utils.clip_grad_norm_(self.optimizer.all_params, max_norm=clip)\n",
    "            for optimizer, scheduler in zip(self.optimizer.optims, self.optimizer.schedulers):\n",
    "                optimizer.step()\n",
    "                scheduler.step()\n",
    "            self.optimizer.zero_grad()\n",
    "\n",
    "            self.step += 1\n",
    "\n",
    "            if batch_idx % log_interval == 0:\n",
    "                elapsed = time.time() - start_time\n",
    "\n",
    "                lrs = self.optimizer.get_lr()\n",
    "                logging.info(\n",
    "                    '| epoch {:3d} | step {:3d} | batch {:3d}/{:3d} | lr{} | loss {:.4f} | s/batch {:.2f}'.format(\n",
    "                        epoch, self.step, batch_idx, self.batch_num, lrs,\n",
    "                        losses / log_interval,\n",
    "                        elapsed / log_interval))\n",
    "\n",
    "                losses = 0\n",
    "                start_time = time.time()\n",
    "\n",
    "            batch_idx += 1\n",
    "\n",
    "        overall_losses /= self.batch_num\n",
    "        during_time = time.time() - epoch_start_time\n",
    "\n",
    "        # reformat\n",
    "        overall_losses = reformat(overall_losses, 4)\n",
    "        score, f1 = get_score(y_true, y_pred)\n",
    "\n",
    "        logging.info(\n",
    "            '| epoch {:3d} | score {} | f1 {} | loss {:.4f} | time {:.2f}'.format(epoch, score, f1,\n",
    "                                                                                  overall_losses,\n",
    "                                                                                  during_time))\n",
    "        if set(y_true) == set(y_pred) and self.report:\n",
    "            report = classification_report(y_true, y_pred, digits=4, target_names=self.target_names)\n",
    "            logging.info('\\n' + report)\n",
    "\n",
    "        return f1\n",
    "\n",
    "    def _eval(self, epoch, test=False):\n",
    "        self.model.eval()\n",
    "        start_time = time.time()\n",
    "\n",
    "        y_pred = []\n",
    "        y_true = []\n",
    "        with torch.no_grad():\n",
    "            for batch_data in data_iter(self.dev_data, test_batch_size, shuffle=False):\n",
    "                torch.cuda.empty_cache()\n",
    "                batch_inputs, batch_labels = self.batch2tensor(batch_data)\n",
    "                batch_outputs = self.model(batch_inputs)\n",
    "                y_pred.extend(torch.max(batch_outputs, dim=1)[1].cpu().numpy().tolist())\n",
    "                y_true.extend(batch_labels.cpu().numpy().tolist())\n",
    "\n",
    "            score, f1 = get_score(y_true, y_pred)\n",
    "\n",
    "            during_time = time.time() - start_time\n",
    "            \n",
    "            if test:\n",
    "                df = pd.DataFrame({'label': y_pred})\n",
    "                df.to_csv(save_test, index=False, sep=',')\n",
    "            else:\n",
    "                logging.info(\n",
    "                    '| epoch {:3d} | dev | score {} | f1 {} | time {:.2f}'.format(epoch, score, f1,\n",
    "                                                                              during_time))\n",
    "                if set(y_true) == set(y_pred) and self.report:\n",
    "                    report = classification_report(y_true, y_pred, digits=4, target_names=self.target_names)\n",
    "                    logging.info('\\n' + report)\n",
    "\n",
    "        return f1\n",
    "\n",
    "    def batch2tensor(self, batch_data):\n",
    "        '''\n",
    "            [[label, doc_len, [[sent_len, [sent_id0, ...], [sent_id1, ...]], ...]]\n",
    "        '''\n",
    "        batch_size = len(batch_data)\n",
    "        doc_labels = []\n",
    "        doc_lens = []\n",
    "        doc_max_sent_len = []\n",
    "        for doc_data in batch_data:\n",
    "            doc_labels.append(doc_data[0])\n",
    "            doc_lens.append(doc_data[1])\n",
    "            sent_lens = [sent_data[0] for sent_data in doc_data[2]]\n",
    "            max_sent_len = max(sent_lens)\n",
    "            doc_max_sent_len.append(max_sent_len)\n",
    "\n",
    "        max_doc_len = max(doc_lens)\n",
    "        max_sent_len = max(doc_max_sent_len)\n",
    "\n",
    "        batch_inputs1 = torch.zeros((batch_size, max_doc_len, max_sent_len), dtype=torch.int64)\n",
    "        batch_inputs2 = torch.zeros((batch_size, max_doc_len, max_sent_len), dtype=torch.int64)\n",
    "        batch_masks = torch.zeros((batch_size, max_doc_len, max_sent_len), dtype=torch.float32)\n",
    "        batch_labels = torch.LongTensor(doc_labels)\n",
    "\n",
    "        for b in range(batch_size):\n",
    "            for sent_idx in range(doc_lens[b]):\n",
    "                sent_data = batch_data[b][2][sent_idx]\n",
    "                for word_idx in range(sent_data[0]):\n",
    "                    batch_inputs1[b, sent_idx, word_idx] = sent_data[1][word_idx]\n",
    "                    batch_inputs2[b, sent_idx, word_idx] = sent_data[2][word_idx]\n",
    "                    batch_masks[b, sent_idx, word_idx] = 1\n",
    "\n",
    "        if use_cuda:\n",
    "            batch_inputs1 = batch_inputs1.to(device)\n",
    "            batch_inputs2 = batch_inputs2.to(device)\n",
    "            batch_masks = batch_masks.to(device)\n",
    "            batch_labels = batch_labels.to(device)\n",
    "\n",
    "        return (batch_inputs1, batch_inputs2, batch_masks), batch_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-07-17 11:38:06,733 INFO: Total 9000 docs.\n",
      "2020-07-17 11:38:10,409 INFO: Total 1000 docs.\n",
      "2020-07-17 11:38:10,410 INFO: Start training...\n"
     ]
    }
   ],
   "source": [
    "# train\n",
    "trainer = Trainer(model, vocab)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# test\n",
    "trainer.test()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
